import os
import os.path as osp
import cv2
import SimpleITK as sitk
import numpy as np
import pandas as pd
from concurrent.futures import ThreadPoolExecutor

root_dir = "/mnt/d/projects/2025医学分类/吴远安数据集20250113/dataDemo4LY_20250112"

def save_gray_image(arr_2d, output_path):
    # 截断值到0~255
    arr_2d = np.clip(arr_2d, 0, 255)
    
    # 归一化到0~1
    arr_2d = arr_2d / 255.0
    
    # 转换为uint8类型
    arr_2d = (arr_2d * 255).astype(np.uint8)
    
    # 保存为灰度图像
    cv2.imwrite(output_path, arr_2d)

def single_test():
    image_meta_data = pd.read_csv(osp.join(root_dir, 'exep_20250112xxx_idMappingAndLabels_0116V3.csv'))
    ids = '10219568'

    # import ipdb; ipdb.set_trace()
    image_name = str(image_meta_data[image_meta_data['ids'] == ids]['phase_A_file_name'].iloc[0])
    max_area_index = int(image_meta_data[image_meta_data['ids'] == ids]['maxAreaIndex'].iloc[0])

    image_path = os.path.join(root_dir, 'pic/exep_20250112xxx/Phase_A', image_name)
    # 读取 nii.gz 文件
    image = sitk.ReadImage(image_path)
    array = sitk.GetArrayFromImage(image)
    arr_2d = array[max_area_index] #取出对应病人最大器官横截面积切片的图像array
    # 保存图片
    # 调用函数保存灰度图像
    save_gray_image(arr_2d, 'test.jpg')

def process_image(infos):
    try:
        idx, irow = infos
        if idx % 100 == 0:
            print(f"Processing image {idx}")

        ids, image_name, max_area_index = irow["ids"], irow["phase_A_file_name"], int(irow["maxAreaIndex"])
        image_path = osp.join(root_dir, 'pic/exep_20250112xxx/Phase_A', image_name)
        # 读取 nii.gz 文件
        image = sitk.ReadImage(image_path)
        array = sitk.GetArrayFromImage(image)
        arr_2d = array[max_area_index] #取出对应病人最大器官横截面积切片的图像array
        output_path = osp.join(root_dir, 'select_2d', f"{image_name}.jpg")
        save_gray_image(arr_2d, output_path)
        return output_path
    except Exception as e:
        print(f"Error processing image {e}")
        import ipdb; ipdb.set_trace()
        return ""

def select_and_save_2d():
    image_meta_data = pd.read_csv(osp.join(root_dir, 'exep_20250112xxx_idMappingAndLabels_0116V3.csv'))
    print(f"image_meta_data.shape: {image_meta_data.shape}")
    
    # 创建输出文件夹
    output_folder = osp.join(root_dir, 'select_2d')
    os.makedirs(output_folder, exist_ok=True)
    
    # 使用多线程处理图像
    results = []
    with ThreadPoolExecutor(max_workers=8) as executor:
       results.extend(executor.map(process_image, image_meta_data.iterrows()))
    results = [x for x in results if x]
    print(f"Processed {len(results)} images")

if __name__ == "__main__":
    select_and_save_2d()